import json
import os
from datetime import datetime
from functools import lru_cache
from django.apps import apps
from django.db import transaction
from dotenv import load_dotenv
from channels.generic.websocket import WebsocketConsumer
import base64
from PIL import Image
import io
import requests
from openai import OpenAI
import oss2
import uuid
import imghdr

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from langchain.schema import Document
import mongoengine as me

# 配置加载
load_dotenv()
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")

# 共享资源初始化
llm = ChatOpenAI(
    model="qwen-plus",
    api_key=DASHSCOPE_API_KEY,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0.7,
    max_tokens=2048
)

embeddings = DashScopeEmbeddings(
    dashscope_api_key=DASHSCOPE_API_KEY,
    model="text-embedding-v1"
)

# 提示模板
personal_prompt = ChatPromptTemplate.from_messages([
    HumanMessagePromptTemplate.from_template("""
你现在是一名专业的AI健身教练，擅长提供科学的健身指导、营养建议和健康管理方案。
你的回答需要基于用户的具体情况，提供可操作的建议。

【用户基本信息】
性别: {gender}
生日: {birthdate}
身高: {height} cm
体重: {weight} kg
BMI: {bmi}

【回答要求】
1. 优先使用检索到的健康知识库内容进行回答，确保回答的专业性和准确性
2. 如果检索结果不充分，请结合你的专业知识提供个性化建议
3. 对于健身相关问题，需考虑用户可能的运动水平（初学者/中级/高级）
4. 提供具体的行动步骤或训练计划时，需明确指出适用人群
5. 遇到健康风险相关问题，建议用户咨询专业人士

【检索结果】
{context}

【历史对话】
{chat_history}

用户: {question}
教练:""")
])

default_prompt = ChatPromptTemplate.from_messages([
    HumanMessagePromptTemplate.from_template("""
你现在是一名专业的AI健身教练，擅长提供科学的健身指导、营养建议和健康管理方案。
你的回答需要基于用户的具体情况，提供可操作的建议。

【回答要求】
1. 优先使用检索到的健康知识库内容进行回答，确保回答的专业性和准确性
2. 如果检索结果不充分，请结合你的专业知识提供个性化建议
3. 对于健身相关问题，需考虑用户可能的运动水平（初学者/中级/高级）
4. 提供具体的行动步骤或训练计划时，需明确指出适用人群
5. 遇到健康风险相关问题，建议用户咨询专业人士

【检索结果】
{context}

【历史对话】
{chat_history}

用户: {question}
教练:""")
])


# 数据加载与缓存优化
@lru_cache(maxsize=1)
def get_vectorstore():
    """持久化健康知識庫向量索引，優先從本地加載"""
    from langchain_community.vectorstores import FAISS
    import os

    # 設定存儲路徑
    index_dir = r"D:\\Work\\Python AI\\P5\\P5Code\\faiss_index"
    index_path = os.path.join(index_dir, "index.faiss")
    pkl_path = os.path.join(index_dir, "index.pkl")

    # 若本地已存在，直接加載
    if os.path.exists(index_path) and os.path.exists(pkl_path):
        return FAISS.load_local(index_dir, embeddings, allow_dangerous_deserialization=True)

    # 否則從資料庫構建並保存
    HealthFAQ = apps.get_model('coach', 'HealthFAQ')
    documents = []
    for faq in HealthFAQ.objects.all():
        documents.append(Document(
            page_content=f"問題：{faq.question}\n回答：{faq.answer}",
            metadata={"question": faq.question, "answer": faq.answer}
        ))
    if documents:
        vectorstore = FAISS.from_documents(documents, embeddings)
        os.makedirs(index_dir, exist_ok=True)
        vectorstore.save_local(index_dir)
        return vectorstore
    return None


# 同步健康数据初始化
def init_health_data():
    """同步初始化健康知识库，确保只执行一次"""
    HealthFAQ = apps.get_model('coach', 'HealthFAQ')

    # 检查是否已初始化
    if HealthFAQ.objects.exists():
        return

    # 使用同步事务
    with transaction.atomic():
        initial_data = [
            {'question': '每天应该喝多少水？',
             'answer': '一般建议成年人每天饮用1500-2000毫升水，运动后可适当增加，避免一次性大量饮用。'},
            {'question': '什么是健康的BMI范围？',
             'answer': 'BMI 18.5-23.9被认为是健康范围，低于18.5为偏瘦，24-27.9为超重，28及以上为肥胖。'},
            {'question': '如何提高睡眠质量？',
             'answer': '保持规律作息、睡前1小时避免蓝光设备、保持卧室安静黑暗、避免睡前摄入咖啡因和大量食物。'}
        ]
        # 批量创建数据
        HealthFAQ.objects.bulk_create(
            [HealthFAQ(**item) for item in initial_data]
        )


# 同步处理函数
def process_question(state: dict) -> dict:
    """同步处理用户问题，检索相关文档"""
    user_input = state["user_input"]
    state["current_question"] = user_input
    state["chat_history"].append({"user": user_input})

    # 获取向量库（已缓存）
    vectorstore = get_vectorstore()

    if vectorstore and user_input:
        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
        docs = retriever.get_relevant_documents(user_input)
        state["source_docs"] = [
            {"question": doc.metadata.get("question"), "answer": doc.metadata.get("answer")}
            for doc in docs
        ]

    return state


def generate_response(state: dict) -> dict:
    """同步生成回答，调用LLM"""
    context = "\n\n".join([f"{d['question']}\n{d['answer']}" for d in state["source_docs"]])
    chat_history_text = "\n".join([
        f"用户: {m['user']}" if 'user' in m else f"教练: {m['coach']}"
        for m in state["chat_history"][:-1]
    ])

    user_profile = state.get("user_profile", {})
    has_profile = any(v not in [None, '', '--'] for v in user_profile.values())

    # 选择对应的提示模板
    if has_profile:
        chain = LLMChain(llm=llm, prompt=personal_prompt)
        response = chain.run(
            context=context,
            chat_history=chat_history_text,
            question=state["current_question"],
            gender=user_profile.get("gender", "--"),
            birthdate=user_profile.get("birthdate", "--"),
            height=user_profile.get("height", "--"),
            weight=user_profile.get("weight", "--"),
            bmi=user_profile.get("bmi", "--"),
        )
    else:
        chain = LLMChain(llm=llm, prompt=default_prompt)
        response = chain.run(
            context=context,
            chat_history=chat_history_text,
            question=state["current_question"]
        )

    state["response"] = response
    state["chat_history"].append({"coach": response})
    return state


# 状态图构建（同步模式）
graph_builder = StateGraph(dict)
graph_builder.add_node("process_question", process_question)
graph_builder.add_node("generate_response", generate_response)
graph_builder.set_entry_point("process_question")
graph_builder.add_edge("process_question", "generate_response")
graph_builder.add_edge("generate_response", END)

memory_saver = MemorySaver()
coach_chain = graph_builder.compile(checkpointer=memory_saver)

# 线程安全的用户会话存储
THREAD_PROFILE_MAP = {}


def upload_to_oss(image_base64, filename):
    OSS_ACCESS_KEY_ID = os.getenv("OSS_ACCESS_KEY_ID")
    OSS_ACCESS_KEY_SECRET = os.getenv("OSS_ACCESS_KEY_SECRET")
    OSS_ENDPOINT = os.getenv("OSS_ENDPOINT")
    OSS_BUCKET_NAME = os.getenv("OSS_BUCKET_NAME")
    auth = oss2.Auth(OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET)
    bucket = oss2.Bucket(auth, OSS_ENDPOINT, OSS_BUCKET_NAME)
    image_bytes = base64.b64decode(image_base64)
    bucket.put_object(filename, image_bytes)
    # 拼接公网URL
    endpoint_netloc = OSS_ENDPOINT.replace('https://', '')
    return f"https://{OSS_BUCKET_NAME}.{endpoint_netloc}/{filename}"


def get_image_extension(image_bytes):
    ext = imghdr.what(None, h=image_bytes)
    if ext == 'jpeg':
        ext = 'jpg'
    return ext or 'png'


class CoachMessage(me.Document):
    user_id = me.StringField()
    thread_id = me.StringField()
    role = me.StringField(choices=['user', 'coach'])
    text = me.StringField()
    image = me.StringField()  # base64或图片url
    timestamp = me.DateTimeField()


def save_message(user_id, thread_id, role, text, image=None):
    CoachMessage(
        user_id=user_id,
        thread_id=thread_id,
        role=role,
        text=text,
        image=image,
        timestamp=datetime.utcnow()
    ).save()


def get_history(user_id, thread_id):
    msgs = CoachMessage.objects(user_id=user_id, thread_id=thread_id).order_by('timestamp')
    return [
        {
            'role': msg.role,
            'text': msg.text,
            'image': msg.image,
            'timestamp': msg.timestamp.isoformat()
        }
        for msg in msgs
    ]


class CoachConsumer(WebsocketConsumer):
    def connect(self):
        self.accept()

    def disconnect(self, close_code):
        pass

    def receive(self, text_data):
        try:
            data = json.loads(text_data)
            msg_type = data.get('type', 'text')
            user_id = data.get('user_id') or (data.get('user_profile', {}) or {}).get('user_id') or 'anonymous'
            thread_id = data.get('thread_id', 'default_thread')

            # 1. 获取历史消息
            if msg_type == 'history':
                history = get_history(user_id, thread_id)
                self.send(text_data=json.dumps({
                    'type': 'history',
                    'history': history
                }, ensure_ascii=False))
                return

            if msg_type == 'clear_history':
                # 删除该用户该线程的所有消息
                CoachMessage.objects(user_id=user_id, thread_id=thread_id).delete()
                self.send(text_data=json.dumps({
                    'type': 'clear_history_success'
                }, ensure_ascii=False))
                return

            # 2. 图片消息
            if msg_type == 'image':
                image_base64 = data.get('image_base64')
                if not image_base64:
                    self.send(text_data=json.dumps({
                        "type": "error",
                        "message": "未收到图片数据"
                    }, ensure_ascii=False))
                    return
                try:
                    # 自动识别图片类型
                    image_bytes = base64.b64decode(image_base64)
                    ext = get_image_extension(image_bytes)
                    filename = f"aiimg/{uuid.uuid4().hex}.{ext}"
                    image_url = upload_to_oss(image_base64, filename)

                    # 用OpenAI SDK调用Qwen-VL
                    client = OpenAI(
                        api_key=os.getenv("DASHSCOPE_API_KEY"),
                        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                    )
                    completion = client.chat.completions.create(
                        model="qwen-vl-plus",
                        messages=[
                            {"role": "user", "content": [
                                {"type": "image_url", "image_url": {"url": image_url}},
                                {"type": "text", "text": """
                                    1.内容描述：用中文简洁、准确地描述图片呈现的核心场景、主体元素及整体氛围；
                                    2.文字识别：完整提取图片中所有可见文字（包括印刷体、手写体等，若有模糊或残缺处请标注）。

                                    额外说明：
                                        若图片主题为跌打扭伤相关（如受伤部位、治疗场景、药品器械等），请补充：
                                        日常预防跌打扭伤的实用建议（如运动前准备、动作规范等）；
                                        不同阶段的科学治疗方法（如急性期处理、恢复期护理等）。
                                        若图片涉及运动健身或户外活动（如跑步、器械锻炼、登山、露营等场景），请补充：
                                        适合该场景的运动方案（如入门级动作、进阶训练计划等）；
                                        运动中的安全注意事项（如装备选择、强度控制、环境适应等）；
                                        提升运动效果的小技巧（如呼吸方法、动作发力要点等）。
                                    注：如果额外说明的场景未体现，则不做运动相关阐述，否则会受到惩罚。
                                """}
                            ]}
                        ]
                    )
                    image_caption = completion.choices[0].message.content if completion.choices else ""
                    # 保存图片消息
                    save_message(user_id, thread_id, 'user', '[图片]', image=image_base64)
                    save_message(user_id, thread_id, 'coach', image_caption)
                    self.send(text_data=json.dumps({
                        "type": "image_result",
                        "image_caption": image_caption
                    }, ensure_ascii=False))
                except Exception as e:
                    self.send(text_data=json.dumps({
                        "type": "error",
                        "message": f"图片识别失败: {str(e)}"
                    }, ensure_ascii=False))
                return

            # 3. 文本消息
            user_input = data.get('user_input', '')
            user_profile = data.get('user_profile', {})
            # 構造檢索上下文
            initial_state = {
                "chat_history": [],
                "current_question": "",
                "user_input": user_input,
                "user_profile": user_profile,
                "response": "",
                "source_docs": []
            }
            # 先檢索知識庫
            state = process_question(initial_state)
            context = "\n\n".join([f"{d['question']}\n{d['answer']}" for d in state["source_docs"]])
            chat_history_text = "\n".join([
                f"用户: {m['user']}" if 'user' in m else f"教练: {m['coach']}"
                for m in state["chat_history"][:-1]
            ])
            user_profile = state.get("user_profile", {})
            has_profile = any(v not in [None, '', '--'] for v in user_profile.values())
            # 構造 prompt
            if has_profile:
                messages = [
                    {"role": "system", "content": "你现在是一名专业的AI健身教练，擅长提供科学的健身指导、营养建议和健康管理方案。你的回答需要基于用户的具体情况，提供可操作的建议。"},
                    {"role": "user", "content": f"【用户基本信息】\n性别: {user_profile.get('gender', '--')}\n生日: {user_profile.get('birthdate', '--')}\n身高: {user_profile.get('height', '--')} cm\n体重: {user_profile.get('weight', '--')} kg\nBMI: {user_profile.get('bmi', '--')}\n\n【回答要求】\n1. 优先使用检索到的健康知识库内容进行回答，确保回答的专业性和准确性\n2. 如果检索结果不充分，请结合你的专业知识提供个性化建议\n3. 对于健身相关问题，需考虑用户可能的运动水平（初学者/中级/高级）\n4. 提供具体的行动步骤或训练计划时，需明确指出适用人群\n5. 遇到健康风险相关问题，建议用户咨询专业人士\n\n【检索结果】\n{context}\n\n【历史对话】\n{chat_history_text}\n\n用户: {user_input}\n教练:"}
                ]
            else:
                messages = [
                    {"role": "system", "content": "你现在是一名专业的AI健身教练，擅长提供科学的健身指导、营养建议和健康管理方案。你的回答需要基于用户的具体情况，提供可操作的建议。"},
                    {"role": "user", "content": f"【回答要求】\n1. 优先使用检索到的健康知识库内容进行回答，确保回答的专业性和准确性\n2. 如果检索结果不充分，请结合你的专业知识提供个性化建议\n3. 对于健身相关问题，需考虑用户可能的运动水平（初学者/中级/高级）\n4. 提供具体的行动步骤或训练计划时，需明确指出适用人群\n5. 遇到健康风险相关问题，建议用户咨询专业人士\n\n【检索结果】\n{context}\n\n【历史对话】\n{chat_history_text}\n\n用户: {user_input}\n教练:"}
                ]
            # 用 OpenAI SDK 直接流式生成
            client = OpenAI(
                api_key=os.getenv("DASHSCOPE_API_KEY"),
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            )
            try:
                completion = client.chat.completions.create(
                    model="qwen-plus",
                    messages=messages,
                    stream=True
                )
                full_response = ""
                for chunk in completion:
                    delta = getattr(chunk.choices[0].delta, "content", None)
                    if delta:
                        full_response += delta
                        self.send(text_data=json.dumps({
                            "type": "stream",
                            "delta": delta
                        }, ensure_ascii=False))
                # 保存消息
                save_message(user_id, thread_id, 'user', user_input)
                save_message(user_id, thread_id, 'coach', full_response)
                self.send(text_data=json.dumps({
                    "type": "stream_end",
                    "user_input": user_input,
                    "response": full_response,
                    "sources": state.get("source_docs", []),
                    "conversation_history": state.get("chat_history", []) ,
                    "thread_id": thread_id
                }, ensure_ascii=False))
            except Exception as e:
                self.send(text_data=json.dumps({
                    "type": "error",
                    "message": f"流式输出失败: {str(e)}"
                }, ensure_ascii=False))
            return

        except Exception as e:
            self.send(text_data=json.dumps({
                "error": f"处理错误: {str(e)}"
            }, ensure_ascii=False))